import os
import sys
import cv2
import dlib
import shutil
import corecode
import numpy as np
import matplotlib
import matplotlib.pyplot as plt
np.set_printoptions(threshold=32, edgeitems=10)

#Gets all the faces in an image, as well as their location.
def allFaces(image):

    face_locations = corecode.face_locations(image)

    aR = []

    for face_location in face_locations:
        top, right, bottom, left = face_location
        face_image = image[top:bottom, left:right]
        aR.append([face_image, face_location])
    return aR


#Compares two faces. I didn't write this.
def compare_faces(known_face_encodings, face_encoding_to_check, tolerance=0.5):
    """
    Compare a list of face encodings against a candidate encoding to see if they match.

    :param known_face_encodings: A list of known face encodings
    :param face_encoding_to_check: A single face encoding to compare against the list
    :param tolerance: How much distance between faces to consider it a match. Lower is more strict. 0.6 is typical best performance.
    :return: A list of True/False values indicating which known_face_encodings match the face encoding to check
    """
    return (corecode.face_distance(known_face_encodings, face_encoding_to_check) <= tolerance)

#A 0.x ieterator.
def frange(start, stop, step):
    i = start
    while i < stop:
        yield i
        i += step
 

#Gets the two images.
cwd = os.getcwd()
sys_a = sys.argv[1]
sys_b = sys.argv[2]
cv_a = cv2.imread(sys_a)
cv_b = cv2.imread(sys_b)

#Makes sure that the two images are actually images.
if sys_a is None:
        print 'Image (a) is not an image'
        sys.exit(1)
if sys_b is None:
        print 'Image (b) is not an image'
        sys.exit(1)

if str(type(cv_a)) == "<type 'NoneType'>":
        print 'Image (a) is not an image'
        sys.exit(1)
if str(type(cv_b)) == "<type 'NoneType'>":
        print 'Image (b) is not an image'
        sys.exit(1)

#Gets the images, faces, locations, cropped images, and encodings.
a_image         = corecode.load_image_file(sys_a)
b_image         = corecode.load_image_file(sys_b)
a_face_e        = corecode.face_encodings(a_image)
b_face_e        = corecode.face_encodings(b_image)
a_location      = corecode.face_locations(a_image)
b_location      = corecode.face_locations(b_image)
a_face_image    = (allFaces(a_image))[0][0]
b_face_image    = (allFaces(b_image))[0][0]
a_face_encoding = corecode.face_encodings(a_image)[0]
b_face_encoding = corecode.face_encodings(b_image)[0]

#Checks for how many faces are there.
if len(a_face_e) > 1:
        print 'This script is for testing two faces, in two images, only'
        print 'Image (a) has two faces.'
        print ("\n\n" + str(len(a_face_encoding)))
        sys.exit(1)
if len(b_face_e) > 1:
        print 'This script is for testing two faces, in two images, only'
        print 'Image (b) has two faces.'
        sys.exit(1)

#Checks the faces against different tollerance levels.
allRS = []
for i in frange(0.0,  1.0,  0.05):
    results = compare_faces([a_face_encoding], b_face_encoding, i)
    allRS.append("Same Person: {} @ tolerance level {}".format(results, i))
    print ("Same Person: {} @ tolerance level {}".format(results, i))


#The matplolib section, to display the images and results.
ff, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(13,8))
ff.suptitle('Person Comparison')

#Display the faces
ax1.imshow(a_face_image)
ax1.set_title((sys_a.split('/')[-1]))
ax1.tick_params(labeltop='off', labelright='off', labelleft='off', labelbottom='off')
ax2.imshow(b_face_image)
ax2.set_title((sys_b.split('/')[-1]))
ax2.tick_params(labeltop='off', labelright='off', labelleft='off', labelbottom='off')

#Display the face comparison results.
ax3.plot([0, 0], 'r')
ax3.grid(False)
ax3.tick_params(labeltop='off', labelright='off', labelleft='off', labelbottom='off')
ax3.axis([0, 10, -len(allRS), 0])
for i, s in enumerate(allRS):
    ax3.text(0.1, -i-0.5, s, fontsize=10)


#Saves the image as a png.
plt.savefig('pc-{}-{}.png'.format((sys_a.split('/')[-1]), (sys_b.split('/')[-1])))
plt.show()


